#include "types.h"
#include "riscv.h"
#include "defs.h"

#define UART0 0x10000000L
#define KERNBASE 0x80000000L
#define PHYSTOP 0x80000000L + 128 * 1024 * 1024 //128MB

extern char etext[];  // kernel.ld中定义的内核结束符号

pagetable_t kernel_pagetable;


int mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm);

// 通过虚拟地址和层级获取虚拟页号
int 
get_idx(int level, uint64 va)
{
//  va的39位虚拟地址格式：
//  |123456789|123456789|123456789|000000000000|
//  38      30 29     21 20     12 11          0
//  +--------+--------+--------+------------+
//  |  VPN2  |  VPN1  |  VPN0  |    page    |
//  +--------+--------+--------+------------+
    int index = (va >> (12 + 9 * level)) & 0x1FF;
    return index;
}

pagetable_t
kvmmake(void)
{
    pagetable_t kpgtbl;

    kpgtbl = (pagetable_t)kalloc();
    memset(kpgtbl, 0, PGSIZE);

    // 映射串口位址
    mappages(kpgtbl, UART0, PGSIZE, UART0, PTE_R | PTE_W);

    // 映射内核代码段
    mappages(kpgtbl, KERNBASE, (uint64)etext - KERNBASE, KERNBASE, PTE_R | PTE_X);

    // 映射内核数据段
    mappages(kpgtbl, (uint64)etext, PHYSTOP - (uint64)etext, (uint64)etext, PTE_R | PTE_W);

    return kpgtbl;
}

void 
kvminit()
{
    kernel_pagetable = kvmmake();
}

void
kvminithart()
{
    w_satp(MAKE_SATP(kernel_pagetable)); // 将页表根地址写入satp寄存器，
                                         // 告诉CPU从该页表开始进行地址转换
    sfence_vma(); // 刷新TLB
}

// 创建一个新的空页表提供给用户进程使用
pagetable_t 
create_pagetable(void) {
    pagetable_t pagetable;

    pagetable = (pagetable_t)kalloc();
    if(pagetable == 0)
        return 0;
    memset(pagetable, 0, PGSIZE);

    return pagetable;
}

// 为已经存在的页表啊分配更多空间
uint64 
grow_pagetable(pagetable_t pagetable, uint64 oldsz, uint64 newsz) {
    if(newsz < oldsz)
        return oldsz;
    
    char *mem;
    uint64 a;

    oldsz = PGROUNDUP(oldsz);
    for(a = oldsz; a < newsz; a += PGSIZE) {
        mem = kalloc();
        if(mem == 0) {
            // 在循环的其他位置可能已经分配了一些页表空间
            // 所以需要将这些空间释放掉，下面同理
            decline_pagetable(pagetable, a, oldsz); 
            return 0;
        }

        memset(mem, 0, PGSIZE);
        if(mappages(pagetable, a, PGSIZE, (uint64)mem, PTE_R | PTE_W) != 0) {
            kfree(mem);
            decline_pagetable(pagetable, a, oldsz);
            return 0;
        }
    }

    return newsz;
}

// 为已经存在的页表释放多余空间
uint64 
decline_pagetable(pagetable_t pagetable, uint64 oldsz, uint64 newsz) {
    if(newsz >= oldsz)
        return oldsz;

    if(PGROUNDDOWN(newsz) < PGROUNDDOWN(oldsz)) {
        int npages = (PGROUNDUP(oldsz) - PGROUNDUP(newsz)) / PGSIZE;
        unmap_pages(pagetable, PGROUNDUP(newsz), npages, 1);
    }

    return newsz;
}

pte_t*
walk(pagetable_t pagetable, uint64 va, int alloc)
{
    for(int level = 2; level > 0; level--) {
        int idx = get_idx(level, va);
        pte_t *pte = &pagetable[idx]; // 根据索引找到页表项

        if(*pte & PTE_V) { // 如果页表项有效
            pagetable = (pagetable_t)PTE2PA(*pte); // 从页表中取出下一级页表的地址
        } else {
            if(!alloc || (pagetable = (pagetable_t)kalloc()) == 0)
                return 0;
            memset(pagetable, 0, PGSIZE); // 不可省略，有效位等标志位都需要初始化为0
            *pte = PA2PTE(pagetable) | PTE_V;
        }
    }

    return &pagetable[get_idx(0, va)]; // 返回最终页表项的物理地址
}


// 返回虚拟地址va映射的物理地址
uint64 
walkaddr(pagetable_t pagetable, uint64 va)
{
    pte_t *pte;

    pte = walk(pagetable, va, 0);
    if(pte == 0)
        return 0;
    if((*pte & PTE_V) == 0)
        return 0;
    if((*pte & PTE_U) == 0)
        return 0;
    return PTE2PA(*pte);
}

// 释放一个页表及其所有下级页表的空间
void
freewalk(pagetable_t pagetable) {
    // 每个页表有512个页表项
    for(int i = 0; i < 512; i++) {
        pte_t pte = pagetable[i];

        // 根据SV39规定，有效且不可读写执行的页表项指向下一级页表
        if((pte & PTE_V) && ((pte & (PTE_R | PTE_W | PTE_X)) == 0)) {
            uint64 child = PTE2PA(pte);
            freewalk((pagetable_t)child);
            pagetable[i] = 0;
        }
        else if(pte & PTE_V) {
            return; // 调用freewalk前，应确保所有页表项都已经被unmap_pages释放
        }
    }

    kfree((void*)pagetable);
}

// 释放整个用户空间的页表及其所有物理内存
void
free_user_pages(pagetable_t pagetable, uint64 sz) {
    if(sz > 0) { // 为什么这里需要sz参数？ // sz似乎是一个需要被用户空间维护的参数
        unmap_pages(pagetable, 0, PGROUNDUP(sz) / PGSIZE, 1);
    }

    freewalk(pagetable);
}

// pte页表项地址格式：
// |000000000000000000000000|123456789|123456789|123456789|
//  63                    10 9       1 0
//  +------------------------+---------+---------+---------+
//  |      物理页            |  标示位  |  标示位  |  标示位
//  +------------------------+---------+---------+---------+
int 
mappages(pagetable_t pagetable, uint64 va, uint64 size, uint64 pa, int perm)
{
    uint64 a, end;
    pte_t *pte;

    if(va % PGSIZE != 0 || size % PGSIZE != 0)
        return -1;

    a = va;
    end = va + size - PGSIZE;

    for(;;) {
        //这里需要给从va到va+size的每一页都映射到pa到pa+size
        if((pte = walk(pagetable, a, 1)) == 0) // 如果找不到对应的页表项
            return -1;
        if(*pte & PTE_V) // 如果页表项已经存在
            return -1;
        *pte = PA2PTE(pa) | perm | PTE_V; // 将物理地址映射为页表项并写入有效位
        if(a == end)
            break;
        a += PGSIZE;
        pa += PGSIZE; //向后分配
    }

    return 0;
}

void
unmap_pages(pagetable_t pagetable, uint64 va, uint64 npages, int do_free) {
    uint64 a;
    pte_t *pte;

    if(va % PGSIZE != 0)
        return;

    for(a = va; a < va + npages * PGSIZE; a += PGSIZE) {
        if((pte = walk(pagetable, a, 0)) == 0) // 若该页表项不存在，则略过
            continue;
        if((*pte & PTE_V) == 0) // 若该页表项无效，则略过
            continue;
        if(do_free) {
            uint64 pa = PTE2PA(*pte);
            kfree((void*)pa);
        }
        *pte = 0;
    }
}

int map_page(pagetable_t pagetable, uint64 va, uint64 pa, int perm) {
    return mappages(pagetable, va, PGSIZE, pa, perm);
}